"""
根据演示demo(characterglm_api_demo_streamlit)改造生成，可以设定两个或者多个角色的人设，通过选择角色交替生成对话。
用户可以设置角色，输入角色名和角色人设，点击【生成图片】按钮，生成一张角色的图片。
用户可以设置用户的名字和用户的人设。
用户可以选择参与对话的角色，在对话输入框中输入对话内容，点击【发送】按钮，生成对话内容。
调用CharacterGLM实现角色扮演，调用CogView生成图片，调用ChatGLM生成CogView所需的prompt。

依赖：
pyjwt
requests
streamlit
zhipuai
python-dotenv

运行方式：
```bash
streamlit run characterglm_chat_streamlit.py
```
"""

import os
import itertools
from typing import Iterator, Optional
import json
import streamlit.components.v1 as components

import streamlit as st
from dotenv import load_dotenv
# 通过.env文件设置环境变量
# reference: https://github.com/theskumar/python-dotenv
load_dotenv()

import api
from api import generate_chat_scene_prompt, generate_role_appearance, get_characterglm_response, generate_cogview_image
from data_types import TextMsg, ImageMsg, TextMsgList, MsgList, filter_text_msg

st.set_page_config(page_title="CharacterGLM API Demo", page_icon="🤖", layout="wide")
debug = os.getenv("DEBUG", "").lower() in ("1", "yes", "y", "true", "t", "on")


def update_api_key(key: Optional[str] = None):
    if debug:
        print(f'update_api_key. st.session_state["API_KEY"] = {st.session_state["API_KEY"]}, key = {key}')
    key = key or st.session_state["API_KEY"]
    if key:
        api.API_KEY = key

# 设置API KEY
api_key = st.sidebar.text_input("API_KEY", value=os.getenv("API_KEY", ""), key="API_KEY", type="password", on_change=update_api_key)
update_api_key(api_key)

# 初始化
if "history" not in st.session_state:
    st.session_state["history"] = []
if "role_list" not in st.session_state:
    st.session_state["role_list"] = []
if "user_meta" not in st.session_state:
    st.session_state["user_meta"] = {
        "user_info": "",
        "user_name": "用户"
    }
if "selected_role_indices" not in st.session_state:
    st.session_state["selected_role_indices"] = []
# 初始化 show_modal
if "show_modal" not in st.session_state:
    st.session_state.show_modal = False

def init_session():
    st.session_state["history"] = []
    st.session_state["role_list"] = []
    st.session_state["user_meta"] = {
        "user_info": "",
        "user_name": "用户"
    }
    st.session_state["selected_role_indices"] = []

# 角色设置区域
with st.expander("角色设置", expanded=True):
    # 让用户输入角色数量
    num_roles = st.number_input("请输入角色数量", min_value=1, value=2, step=1)

    # 调整 meta_list 的长度以匹配用户输入的角色数量
    while len(st.session_state["role_list"]) < num_roles:
        st.session_state["role_list"].append({
            "bot_info": "",
            "bot_name": "",
            "image_urls": []
        })
    while len(st.session_state["role_list"]) > num_roles:
        st.session_state["role_list"].pop()

    # 创建角色选项卡
    tab_names = [f"角色 {i + 1}" for i in range(num_roles)]
    tabs = st.tabs(tab_names)

    for i, tab in enumerate(tabs):
        with tab:
            st.header(f"创建角色 {i + 1}")
            bot_name = st.text_input(f"角色名 {i + 1}", key=f"bot_name_{i}", value=st.session_state["role_list"][i]["bot_name"])
            bot_info = st.text_area(f"角色人设 {i + 1}", key=f"bot_info_{i}", value=st.session_state["role_list"][i]["bot_info"])

            # 以网格形式展示图片
            image_urls = st.session_state["role_list"][i]["image_urls"]
            n_images = len(image_urls)
            if n_images > 0:
                cols_per_row = 3  # 每行显示3张图片
                for j in range(0, n_images, cols_per_row):
                    cols = st.columns(cols_per_row)
                    for k in range(cols_per_row):
                        if j + k < n_images:
                            cols[k].image(image_urls[j + k], caption=f"{bot_name} 的图片")

            if st.button(f"保存角色 {i + 1}", help="点击此按钮保存角色的信息数据。"):
                if bot_name and bot_info:
                    st.session_state["role_list"][i] = {
                        "bot_info": bot_info,
                        "bot_name": bot_name,
                        "image_urls": st.session_state["role_list"][i]["image_urls"] 
                    }
                    st.rerun()

# 用户信息区域
with st.expander("用户信息", expanded=True):
    st.text_input(label="用户名", value=st.session_state["user_meta"]["user_name"], key="user_name", on_change=lambda : st.session_state["user_meta"].update(user_name=st.session_state["user_name"]), help="用户的名字，默认为用户")
    st.text_area(label="用户人设", value=st.session_state["user_meta"]["user_info"], key="user_info", on_change=lambda : st.session_state["user_meta"].update(user_info=st.session_state["user_info"]), help="用户的详细人设信息，可以为空")

# 选择要参与对话的角色
if st.session_state["role_list"]:
    role_names = [role["bot_name"] for role in st.session_state["role_list"]]
    selected_role_names = st.multiselect("选择参与对话的角色", role_names)
    st.session_state["selected_role_indices"] = [role_names.index(name) for name in selected_role_names]

    # 检查是否有未选择的角色
    if len(st.session_state["selected_role_indices"]) < len(st.session_state["role_list"]):
        st.info("检测到有新建角色未被选择，若要为新建角色生成图片，请在上方列表中选择该角色。")

def verify_meta(role_index) -> bool:
    if role_index == -1:
        st.error("请先选择对话角色")
        return False
    selected_role = st.session_state["role_list"][role_index]
    if selected_role["bot_name"] == "" or selected_role["bot_info"] == "":
        st.error(f"所选角色 {selected_role['bot_name']} 的角色名和角色人设不能为空")
        return False
    return True

def draw_new_image():
    # 注入自定义 CSS
    st.markdown(
        """
        <style>
            div[data-testid="stAppViewContainer"] .custom-modal {
                background-color: #f8f9fa;
                border-radius: 10px;
                padding: 20px;
                box-shadow: 0 4px 8px 0 rgba(0, 0, 0, 0.2);
                animation: fadeIn 0.5s;
            }
            @keyframes fadeIn {
                from { opacity: 0; }
                to { opacity: 1; }
            }
            div[data-testid="stAppViewContainer"] .custom-modal h3 {
                color: #007bff;
                font-size: 24px;
                margin-bottom: 20px;
            }
            div[data-testid="stAppViewContainer"] .custom-modal select {
                width: 100%;
                padding: 8px;
                margin-bottom: 20px;
                border: 1px solid #ced4da;
                border-radius: 4px;
            }
            div[data-testid="stAppViewContainer"] .custom-modal button {
                background-color: #007bff;
                color: white;
                border: none;
                border-radius: 5px;
                padding: 8px 16px;
                cursor: pointer;
                transition: background-color 0.3s;
            }
            div[data-testid="stAppViewContainer"] .custom-modal button:hover {
                background-color: #0056b3;
            }
            div[data-testid="stAppViewContainer"] .custom-modal button:active {
                background-color: #003d80;
            }
        </style>
        """,
        unsafe_allow_html=True
    )

    # 模态框检查
    if not st.session_state.show_modal:
        return

    """生成一张图片，并展示在页面上"""
    # 暂时假设使用第一个选择的角色
    # 检查是否有选择的角色
    if len(st.session_state["selected_role_indices"]) == 0:
        if len(st.session_state["role_list"]) == 0:
            st.error("请先创建角色并选择参与对话的角色。")
        else:
            st.error("请先在上方列表中选择要生成图片的角色。")
        return

    selected_role_names = [st.session_state["role_list"][i]["bot_name"] for i in st.session_state["selected_role_indices"]]
    # 使用 st.modal 显示确认提示
    if "show_modal" not in st.session_state:
        st.session_state.show_modal = False

    if st.session_state.show_modal:
        modal_placeholder = st.empty()
        with modal_placeholder.container():
            st.markdown('<div class="custom-modal">', unsafe_allow_html=True)
            st.markdown("### 确认生成图片")
            st.markdown("请选择图片的风格，然后点击【确认生成】按钮开始生成图片。")

            style_options = ["二次元风格", "写实风格", "卡通风格"]
            # 使用单列布局选择风格
            selected_style = st.selectbox("选择图片风格", style_options)        

            st.write(f"您即将为以下角色生成图片：{', '.join(selected_role_names)}。确认要继续吗？")
            st.markdown("**注意：** 生成图片可能需要一些时间，请耐心等待。")

            # 使用两列布局排列按钮
            col1, col2 = st.columns(2)
            with col1:
                st.button("确认生成", key="confirm_generate", on_click=lambda: handle_confirm(selected_role_names, selected_style))
            with col2:
                st.button("取消", key="cancel_generate", on_click=lambda: handle_cancel())
            
            st.markdown('</div>', unsafe_allow_html=True)

# 新增处理函数
def handle_confirm(selected_role_names, selected_style):
    st.session_state.show_modal = False
    # 复制原有的确认生成逻辑到这里
    for index in st.session_state["selected_role_indices"]:
        if not verify_meta(index):
            continue
        
        selected_role = st.session_state["role_list"][index]
        text_messages = filter_text_msg(st.session_state["history"])

        if text_messages:
            image_prompt = "".join(
                generate_chat_scene_prompt(
                    text_messages[-10: ],
                    meta={
                        "bot_info": selected_role["bot_info"],
                        "bot_name": selected_role["bot_name"],
                        "user_info": st.session_state["user_meta"]["user_info"],
                        "user_name": st.session_state["user_meta"]["user_name"]
                    }
                )
            )
        else:
            image_prompt = "".join(generate_role_appearance(selected_role["bot_info"]))
        
        if not image_prompt:
            st.error("调用chatglm生成Cogview prompt出错")
            continue
        
        image_prompt = selected_style + "。" + image_prompt.strip()
        
        if debug:
            print(f"image_prompt = {image_prompt}")
        n_retry = 3
        # st.markdown("正在生成图片，请稍等...")
        success = False
        for i in range(n_retry):
            try:
                img_url = generate_cogview_image(image_prompt)
                if img_url:
                    success = True
            except Exception as e:
                if debug:
                    print(f"生成图片时出错，第 {i + 1} 次重试，错误信息: {e}")
                if i < n_retry - 1:
                    st.error("遇到了一点小问题，重试中...")
                else:
                    st.error("又失败啦，点击【生成图片】按钮可再次重试")
                    continue
            else:
                break
        
        max_images = 5
        st.session_state["role_list"][index]["image_urls"] = (st.session_state["role_list"][index]["image_urls"] + [img_url])[-max_images:]
        
        img_msg = ImageMsg({"role": "image", "image": img_url, "caption": image_prompt})
        while st.session_state["history"] and st.session_state["history"][-1]["role"] == "image":
            st.session_state["history"].pop()
        st.session_state["history"].append(img_msg)
        st.session_state["last_selected_roles"] = st.session_state["selected_role_indices"]
        # st.success("图片生成成功！若要为其他角色生成图片，请重新选择角色。")
    
    if success: 
        if debug:
            print("图片生成成功，即将重新运行脚本")
        # st.rerun()

def handle_cancel():
    st.session_state.show_modal = False
    # st.rerun()

button_labels = {
    "clear_meta": "清空人设",
    "clear_history": "清空对话历史",
    "gen_picture": "生成图片",
    "save_history": "保存对话记录" 
}
if debug:
    button_labels.update({
        "show_api_key": "查看API_KEY",
        "show_meta": "查看meta",
        "show_history": "查看历史"
    })

# 在同一行排列按钮
with st.container():
    n_button = len(button_labels)
    cols = st.columns(n_button)
    button_key_to_col = dict(zip(button_labels.keys(), cols))
    
    with button_key_to_col["clear_meta"]:
        clear_meta = st.button(button_labels["clear_meta"], key="clear_meta")
        if clear_meta:
            # 清空角色列表
            st.session_state["role_list"] = []
            # 清空用户人设
            st.session_state["user_meta"] = {
                "user_info": "",
                "user_name": "用户"
            }
            st.rerun()

    with button_key_to_col["clear_history"]:
        clear_history = st.button(button_labels["clear_history"], key="clear_history")
        if clear_history:
            init_session()
            st.rerun()
    
    with button_key_to_col["gen_picture"]:
        if st.button(button_labels["gen_picture"], key="gen_picture"):
            st.session_state.show_modal = True

    with button_key_to_col["save_history"]:
        save_history = st.button(button_labels["save_history"], key="save_history")
        if save_history:
            try:
                # 初始化分类后的对话记录
                categorized_history = {}
                user_messages = []
                for msg in st.session_state["history"]:
                    if isinstance(msg, dict):
                        if msg["role"] == "user":
                            user_messages.append(msg)
                        elif msg["role"] != "image":
                            if msg["role"] not in categorized_history:
                                categorized_history[msg["role"]] = []
                            # 确保用户消息和角色消息交替存储
                            if user_messages:
                                categorized_history[msg["role"]].extend(user_messages)
                                user_messages = []
                            categorized_history[msg["role"]].append(msg)

                # 处理剩余的用户消息
                for role in categorized_history:
                    if user_messages:
                        categorized_history[role].extend(user_messages)

                with open('chat_history.json', 'w', encoding='utf-8') as f:
                    json.dump(categorized_history, f, ensure_ascii=False, indent=4)
                st.success("对话记录已保存到 chat_history.json")
            except Exception as e:
                st.error(f"保存对话记录时出错: {e}")

    if debug:
        with button_key_to_col["show_api_key"]:
            show_api_key = st.button(button_labels["show_api_key"], key="show_api_key")
            if show_api_key:
                print(f"API_KEY = {api.API_KEY}")
        
        with button_key_to_col["show_meta"]:
            show_meta = st.button(button_labels["show_meta"], key="show_meta")
            if show_meta:
                print(f"用户人设 = {st.session_state['user_meta']}")
                print(f"角色列表 = {st.session_state['role_list']}")
        
        with button_key_to_col["show_history"]:
            show_history = st.button(button_labels["show_history"], key="show_history")
            if show_history:
                print(f"history = {st.session_state['history']}")

# 展示对话历史
for msg in st.session_state["history"]:
    if not isinstance(msg, dict) or "role" not in msg:
        st.error(f"消息格式错误: {msg}")
        continue
    if msg["role"] == "user":
        if "content" not in msg:
            st.error(f"用户消息缺少内容: {msg}")
            continue
        with st.chat_message(name="user", avatar="user"):
            st.markdown(f"**{st.session_state['user_meta']['user_name']}:** {msg['content']}")
    elif msg["role"] == "assistant":
        if "content" not in msg:
            st.error(f"助手消息缺少内容: {msg}")
            continue
        with st.chat_message(name="assistant", avatar="assistant"):
            st.markdown(f"**助手:** {msg['content']}")
    elif msg["role"] == "image":
        if "image" not in msg:
            st.error(f"图片消息缺少图片链接: {msg}")
            continue
        # 尝试查找关联的角色名
        role_name = "未知角色"
        for role in st.session_state["role_list"]:
            if role["bot_name"] in [msg.get("role_name", ""), msg["role"]]:
                role_name = role["bot_name"]
                break
        with st.chat_message(name="assistant", avatar="assistant"):
            st.markdown(f"**{role_name}:**")
            st.image(msg["image"], caption=msg.get("caption", None))
    else:
        # 处理角色名称回复
        role_name = msg["role"]
        with st.chat_message(name=msg["role"], avatar="assistant"):
            if "content" in msg:
                st.markdown(f"**{role_name}:** {msg['content']}")
            elif "image" in msg:
                st.markdown(f"**{role_name}:**")
                st.image(msg["image"], caption=msg.get("caption", None))
            else:
                st.error(f"未知消息格式: {msg}")

# if gen_picture:
#     draw_new_image()


with st.chat_message(name="user", avatar="user"):
    input_placeholder = st.empty()
# with st.chat_message(name="assistant", avatar="assistant"):
#     message_placeholder = st.empty()


def output_stream_response(response_stream: Iterator[str], placeholder):
    content = ""
    for content in itertools.accumulate(response_stream):
        placeholder.markdown(content)
    return content

def start_role_chat():
    query = st.chat_input("开始对话吧")
    if not query:
        return
    else:
        if len(st.session_state["selected_role_indices"]) == 0:
            st.error("请先选择参与对话的角色")
            return
        if not api.API_KEY:
            st.error("未设置API_KEY")

        input_placeholder.markdown(query)
        st.session_state["history"].append(TextMsg({"role": "user", "content": query}))
        
        # 让用户选定的角色进行回复
        for index in st.session_state["selected_role_indices"]:
            selected_role = st.session_state["role_list"][index]
            role_name = selected_role["bot_name"]  # 定义 role_name

            if not verify_meta(index):
                continue
            
            # 为每个角色创建独立的聊天消息框和占位符
            with st.chat_message(name=selected_role["bot_name"], avatar="assistant"):
                st.markdown(f"**{role_name}:**")  # 使用定义好的 role_name
                role_placeholder = st.empty()
                response_stream = get_characterglm_response(filter_text_msg(st.session_state["history"]), meta={
                    "bot_info": selected_role["bot_info"],
                    "bot_name": selected_role["bot_name"],
                    "user_info": st.session_state["user_meta"]["user_info"],
                    "user_name": st.session_state["user_meta"]["user_name"]
                })
                bot_response = output_stream_response(response_stream, role_placeholder)
                if not bot_response:
                    role_placeholder.markdown("生成出错")
                    st.session_state["history"].pop()
                else:
                    st.session_state["history"].append(TextMsg({"role": selected_role["bot_name"], "content": bot_response}))

start_role_chat()
# 调用 draw_new_image 方法
draw_new_image()
